from .stereo import STEREO
from .stereo_tfhead import STEREO_TF
from .stereo_transformer import STEREOTrans
import torch.distributed as dist
from pcdet.utils.common_utils import create_logger
# from .stereo_trt import STEREOTransTRT

__all__ = {
    'stereo': STEREO,
    'stereo_tf':STEREO_TF,
    'stereo_trans': STEREOTrans,
    # 'stereo_trans_trt': STEREOTransTRT
}


def build_detector(model_cfg, num_class, dataset, trt=False):
    if trt:
        model_cfg.NAME +='_trt'
        from .stereo_trt import STEREOTransTRT
        __all__['stereo_trans_trt'] = STEREOTransTRT
    model = __all__[model_cfg.NAME](
        model_cfg=model_cfg, num_class=num_class, dataset=dataset
    )

    try:
        logger = create_logger(rank=dist.get_rank())
    except:
        logger = create_logger()
    if hasattr(model_cfg, 'PRETRAINED_MODEL') and model_cfg.PRETRAINED_MODEL:
        model.load_params_from_file(
            filename=model_cfg.PRETRAINED_MODEL, to_cpu=True, logger=logger)

    return model
